# split caltech-101 to three sets
# author-by: xjtu-blacksmith
# create-on: 2020.2.19

import os
from os import path
import shutil

import numpy as np
import progressbar
from sklearn.model_selection import train_test_split

data_dir = path.join('data', 'caltech-101')  # where the raw data are
background_dir = path.join(data_dir, 'BACKGROUND_Google')  # redundancy, should be removed from data
if path.isdir(background_dir):
    shutil.rmtree(background_dir)  # remove the useless category
categories = os.listdir(data_dir)  # get 101 categories

bar = progressbar.ProgressBar()  # set a default progressbar

for i in bar(range(len(categories))):

    category = categories[i]  # fetch a specific category
    cat_dir = path.join(data_dir, category)

    images = os.listdir(cat_dir)  # get all images under the category

    # images split by their names
    images, images_test = train_test_split(images, test_size=0.25)
    images_train, images_val = train_test_split(images, test_size=0.33)
    image_sets = images_train, images_test, images_val
    labels = 'train', 'test', 'val'

    # move to corresponding folders
    for image_set, label in zip(image_sets, labels):
        dst_folder = path.join(data_dir, label, category)  # create folder
        os.makedirs(dst_folder)
        for image in image_set:
            # rename ../cat/xxx.jpg to ../label/cat/xxx.jpg
            src_dir = path.join(cat_dir, image)
            dst_dir = path.join(dst_folder, image)
            os.rename(src_dir, dst_dir)  # move
    
    os.rmdir(cat_dir)  # remove empty folder
